package cn.xmoit.controller;

import okhttp3.MediaType;
import okhttp3.OkHttpClient;
import okhttp3.Request;
import okhttp3.RequestBody;
import okhttp3.Response;
import org.springframework.stereotype.Controller;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestMethod;
import org.springframework.web.bind.annotation.ResponseBody;

import java.io.IOException;
import java.util.concurrent.TimeUnit;

@Controller
@RequestMapping(value = "/chat")
public class ChatGPTController {

	private static final String API_ENDPOINT = "https://api.openai.com/v1/completions";

	private static final String API_KEY = "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";

	@ResponseBody
	@RequestMapping(value = "/get-chat-info", produces = "application/json;charset=UTF-8",
			method = { RequestMethod.GET, RequestMethod.POST })
	public String getChatInfo(String prompt) throws IOException {
		// OkHttpClient client = new OkHttpClient();
		OkHttpClient client = new OkHttpClient.Builder().connectTimeout(180, TimeUnit.SECONDS)
			.readTimeout(180, TimeUnit.SECONDS)
			.build();

		// String prompt = "I like to eat pizza because";

		MediaType mediaType = MediaType.parse("application/json; charset=utf-8");
		String requestBody = "{\"model\": \"text-davinci-003\",\"prompt\":\"" + prompt
				+ "\",\"max_tokens\":1024,\"temperature\":0}";

		Request request = new Request.Builder().url(API_ENDPOINT)
			.header("Authorization", "Bearer " + API_KEY)
			.header("Content-Type", "application/json")
			.post(RequestBody.create(mediaType, requestBody))
			.build();

		okhttp3.Call call = client.newCall(request);
		call.timeout().timeout(180, TimeUnit.SECONDS);
		Response response = call.execute();

		String responseBody = response.body().string();
		System.out.println(responseBody);

		return responseBody;
	}

}
